import os
import subprocess
import numpy as np


def precision_test(filename):

    def read_output_file(output):
        with open(output, 'r') as file:
            data = file.read().strip()
            return np.array([float(x) for x in data.split()])

    subprocess.run(
        ['python3', filename, '--output=npu.output.txt'],
        env={**os.environ, 'TRITON_DEBUG': '1'},
        check=True
    )

    subprocess.run(
        ['python3', filename, '--output=interpret.output.txt'],
        env={**os.environ, 'TRITON_INTERPRET': '1', 'TRITON_DEBUG': '1'},
        check=True
    )

    npu_output = read_output_file('npu.output.txt')
    interpret_output = read_output_file('interpret.output.txt')

    assert len(npu_output) == len(interpret_output)
    atol = 1e-5
    rtol = 1e-5

    close_mask = np.isclose(npu_output, interpret_output, atol=atol, rtol=rtol)

    differences = np.abs(npu_output - interpret_output)
    differences[close_mask] = 0
    max_diff_index = np.argmax(differences)
    max_diff_value = differences[max_diff_index]

    print(f"{filename} Largest difference at index {max_diff_index}: {max_diff_value}")

def precision_test_bwd(test_filename):
    """
    反向核精度测试：比对NPU执行与Interpret解释器的输出梯度
    :param test_filename: 被测测试文件路径（如 'test_chunk_bwd_kernel_dqkwg.py'）
    """

    npu_prefix = "npu_bwd"
    interpret_prefix = "interpret_bwd"
    atol = 1e-5
    rtol = 1e-5
    required_outputs = ["dq", "dk"]
    optional_outputs = {"dg": True, "dw": True}

    def read_output(prefix, output_name):
        file_path = f"{prefix}_{output_name}.output.txt"
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Output file not found: {file_path}")
        with open(file_path, 'r') as f:
            lines = [line.strip() for line in f if line.strip() and not line.startswith('#')]
        return np.array([float(line) for line in lines], dtype=np.float32)

    print("Running NPU mode...")
    subprocess.run(
        [
            'python3', test_filename,
            '--output-prefix', npu_prefix
        ],
        env={**os.environ, 'TRITON_DEBUG': '1'},
        check=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True
    )

    print("Running Interpret mode...")
    subprocess.run(
        [
            'python3', test_filename,
            '--output-prefix', interpret_prefix
        ],
        env={**os.environ, 'TRITON_INTERPRET': '1', 'TRITON_DEBUG': '1'},
        check=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True
    )

    all_pass = True
    for output_name in required_outputs:
        print(f"\n=== Comparing required output: {output_name} ===")
        npu_data = read_output(npu_prefix, output_name)
        interpret_data = read_output(interpret_prefix, output_name)

        assert len(npu_data) == len(interpret_data), \
            f"{output_name} length mismatch: NPU={len(npu_data)}, Interpret={len(interpret_data)}"

        valid_mask = np.isfinite(npu_data) & np.isfinite(interpret_data)
        if not np.any(valid_mask):
            raise ValueError(f"All values of {output_name} are NaN/Inf!")
        npu_valid = npu_data[valid_mask]
        interpret_valid = interpret_data[valid_mask]
        close_mask = np.isclose(npu_valid, interpret_valid, atol=atol, rtol=rtol)
        match_rate = np.mean(close_mask) * 100
        max_diff = np.max(np.abs(npu_valid - interpret_valid[~close_mask])) if np.any(~close_mask) else 0.0
        max_diff_idx = np.argmax(np.abs(npu_valid - interpret_valid)) if len(npu_valid) > 0 else -1

        print(f"Match rate: {match_rate:.2f}%")
        print(f"Largest difference: {max_diff:.8f} (at index {max_diff_idx})")
        print(f"Number of mismatched values: {np.sum(~close_mask)} / {len(close_mask)}")

        if not np.all(close_mask):
            print(f"FAIL: {output_name} has mismatched values!")
            all_pass = False

    for output_name, is_enabled in optional_outputs.items():
        if not is_enabled:
            print(f"\nSkipping optional output: {output_name} (disabled)")
            continue
        print(f"\n=== Comparing optional output: {output_name} ===")
        try:
            npu_data = read_output(npu_prefix, output_name)
            interpret_data = read_output(interpret_prefix, output_name)
        except FileNotFoundError:
            print(f"FAIL: {output_name} is enabled but output file not found!")
            all_pass = False
            continue

        assert len(npu_data) == len(interpret_data), \
            f"{output_name} length mismatch: NPU={len(npu_data)}, Interpret={len(interpret_data)}"

        valid_mask = np.isfinite(npu_data) & np.isfinite(interpret_data)
        if not np.any(valid_mask):
            print(f"FAIL: All values of {output_name} are NaN/Inf!")
            all_pass = False
            continue
        npu_valid = npu_data[valid_mask]
        interpret_valid = interpret_data[valid_mask]
        close_mask = np.isclose(npu_valid, interpret_valid, atol=atol, rtol=rtol)
        match_rate = np.mean(close_mask) * 100
        max_diff = np.max(np.abs(npu_valid - interpret_valid[~close_mask])) if np.any(~close_mask) else 0.0

        print(f"Match rate: {match_rate:.2f}%")
        print(f"Largest difference: {max_diff:.8f}")
        if not np.all(close_mask):
            print(f"FAIL: {output_name} has mismatched values!")
            all_pass = False

    assert all_pass, "Some outputs failed precision check! See details above."
    print("\n=== All outputs passed precision check! ===")
    
'''
Pass
'''
# precision_test('layer_norm_gated_fwd_kernel/test_layer_norm_gated_fwd_kernel.py')
# precision_test('l2norm_fwd_kernel/test_l2norm_fwd_kernel.py')
# precision_test('chunk_local_cumsum_scalar_kernel/test_chunk_local_cumsum_scalar_kernel.py')
# precision_test('chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra/test_chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra.py')
'''
Fail
'''
# precision_test('causal_conv1d_update_kernel/test_causal_conv1d_update_kernel.py')
# precision_test('solve_tril_16x16_kernel/test_solve_tril_16x16_kernel.py')
'''
TODO, compile error
'''
# test_causal_conv1d_fwd_kernel
# test_causal_conv1d_states_fwd_kernel
# test_chunk_fwd_kernel_o
# test_chunk_bwd_kernel_dqkwg
# test_chunk_bwd_kernel_dv_local
# test_chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64
# test_chunk_gated_delta_rule_fwd_kernel_h_blockdim64
# test_chunk_local_cumsum_vector_kernel
# test_chunk_scaled_dot_kkt_fwd_kernel
# test_chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter
# test_fused_recurrent_gated_delta_rule_fwd_kernel
# test_merge_16x16_to_32x32_inverse_kernel
# test_merge_16x16_to_64x64_inverse_kernel
# test_prepare_wy_repr_bwd_kernel
# test_recompute_w_u_fwd_kernel

# precision_test('causal_conv1d_fwd_kernel/test_causal_conv1d_fwd_kernel.py')
# precision_test('causal_conv1d_states_fwd_kernel/test_causal_conv1d_states_fwd_kernel.py')
# precision_test('chunk_fwd_kernel_o/test_chunk_fwd_kernel_o.py')
# precision_test_bwd('chunk_bwd_kernel_dqkwg/test_chunk_bwd_kernel_dqkwg.py') # compile error



